import numpy as np
import matplotlib.pyplot as plt
import random
import math

class Node:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.parent = None

class RectangleObstacle:
    def __init__(self, x, y, width, height, angle=0):
        self.x = x          # 中心点x坐标
        self.y = y          # 中心点y坐标
        self.width = width  # 宽度
        self.height = height# 高度
        self.angle = angle  # 旋转角度（弧度）
        
    def is_inside(self, point_x, point_y):
        # 将点坐标转换到矩形局部坐标系
        dx = point_x - self.x
        dy = point_y - self.y
        
        # 应用旋转变换
        if self.angle != 0:
            cos_angle = math.cos(-self.angle)
            sin_angle = math.sin(-self.angle)
            rotated_dx = dx * cos_angle - dy * sin_angle
            rotated_dy = dx * sin_angle + dy * cos_angle
            dx, dy = rotated_dx, rotated_dy
        
        # 检查是否在矩形范围内
        return abs(dx) <= self.width/2 and abs(dy) <= self.height/2

class RRT:
    def __init__(self, start, goal, obstacle_list, rand_area, expand_dist=10.0, goal_sample_rate=5, max_iter=1000):
        self.start = Node(start[0], start[1])
        self.goal = Node(goal[0], goal[1])
        self.obstacle_list = obstacle_list
        self.min_rand = rand_area[0]
        self.max_rand = rand_area[1]
        self.expand_dist = expand_dist
        self.goal_sample_rate = goal_sample_rate
        self.max_iter = max_iter
        self.node_list = []

    def planning(self, animation=True):
        self.node_list = [self.start]
        for i in range(self.max_iter):
            rnd_node = self.get_random_node()
            nearest_node = self.get_nearest_node(self.node_list, rnd_node)
            new_node = self.steer(nearest_node, rnd_node, self.expand_dist)

            if self.check_collision(new_node, self.obstacle_list):
                self.node_list.append(new_node)

            if animation and i % 10 == 0:
                self.draw_graph(rnd_node)

            if self.calc_dist_to_goal(self.node_list[-1].x, self.node_list[-1].y) <= self.expand_dist:
                final_node = self.steer(self.node_list[-1], self.goal, self.expand_dist)
                if self.check_collision(final_node, self.obstacle_list):
                    return self.generate_final_course(len(self.node_list) - 1)

            if animation and i % 10 == 0:
                self.draw_graph(rnd_node)

        return None  # cannot find path

    def steer(self, from_node, to_node, extend_length=float("inf")):
        new_node = Node(from_node.x, from_node.y)
        d, theta = self.calc_distance_and_angle(new_node, to_node)

        if extend_length > d:
            new_node.x = to_node.x
            new_node.y = to_node.y
        else:
            new_node.x = from_node.x + extend_length * math.cos(theta)
            new_node.y = from_node.y + extend_length * math.sin(theta)

        new_node.parent = from_node
        return new_node

    def generate_final_course(self, goal_ind):
        path = [[self.goal.x, self.goal.y]]
        node = self.node_list[goal_ind]
        while node.parent is not None:
            path.append([node.x, node.y])
            node = node.parent
        path.append([node.x, node.y])

        return path

    def calc_dist_to_goal(self, x, y):
        dx = x - self.goal.x
        dy = y - self.goal.y
        return math.sqrt(dx ** 2 + dy ** 2)

    def get_random_node(self):
        if random.randint(0, 100) > self.goal_sample_rate:
            rnd = Node(random.uniform(self.min_rand, self.max_rand[0]),
                       random.uniform(self.min_rand, self.max_rand[1]))
        else:  # goal point sampling
            rnd = Node(self.goal.x, self.goal.y)
        return rnd

    def draw_graph(self, rnd=None):
        plt.clf()
        # for stopping simulation with the esc key.
        plt.gcf().canvas.mpl_connect(
            'key_release_event',
            lambda event: [exit(0) if event.key == 'escape' else None])
        if rnd is not None:
            plt.plot(rnd.x, rnd.y, "^k")
            
        for node in self.node_list:
            if node.parent:
                plt.plot([node.x, node.parent.x], [node.y, node.parent.y], "-g")

        # 绘制矩形障碍物
        for obstacle in self.obstacle_list:
            self.plot_rectangle(obstacle)

        plt.plot(self.start.x, self.start.y, "xr")
        plt.plot(self.goal.x, self.goal.y, "xr")
        plt.axis("equal")
        plt.axis([-10, 650, -10, 490])  # 调整坐标轴范围以适应640x480的图像
        plt.grid(True)
        plt.pause(0.01)

    @staticmethod
    def plot_rectangle(obstacle):
        # 获取矩形四个顶点
        w_2 = obstacle.width / 2
        h_2 = obstacle.height / 2
        
        # 四个顶点在局部坐标系中的坐标
        points_local = [
            (w_2, h_2),
            (-w_2, h_2),
            (-w_2, -h_2),
            (w_2, -h_2),
            (w_2, h_2)  # 回到第一个点以闭合多边形
        ]
        
        # 应用旋转变换
        if obstacle.angle != 0:
            cos_angle = math.cos(obstacle.angle)
            sin_angle = math.sin(obstacle.angle)
            points_rotated = []
            for x_local, y_local in points_local:
                x_rotated = x_local * cos_angle - y_local * sin_angle
                y_rotated = x_local * sin_angle + y_local * cos_angle
                points_rotated.append((x_rotated, y_rotated))
            points_local = points_rotated
        
        # 转换到全局坐标系
        points_global = [(x_local + obstacle.x, y_local + obstacle.y) for x_local, y_local in points_local]
        
        # 绘制矩形
        plt.plot([p[0] for p in points_global], [p[1] for p in points_global], 'b-')

    @staticmethod
    def get_nearest_node(node_list, rnd_node):
        dlist = [(node.x - rnd_node.x) ** 2 + (node.y - rnd_node.y) ** 2
                 for node in node_list]
        minind = dlist.index(min(dlist))
        return node_list[minind]

    def check_collision(self, node, obstacle_list):
        # 检查点是否在障碍物内
        for obstacle in obstacle_list:
            if obstacle.is_inside(node.x, node.y):
                return False  # collision
        
        # 检查从父节点到当前节点的线段是否穿过障碍物
        if node.parent:
            return self.check_line_collision(node.parent.x, node.parent.y, node.x, node.y, obstacle_list)
        return True  # 没有父节点时只检查点是否在障碍物内
    
    def check_line_collision(self, x1, y1, x2, y2, obstacle_list):
        # 检查线段是否穿过任何障碍物
        for obstacle in obstacle_list:
            # 简化检查：检查线段的两个端点是否在障碍物内
            if obstacle.is_inside(x1, y1) or obstacle.is_inside(x2, y2):
                return False
                
            # 更复杂的线段与矩形相交检查
            # 这里使用分离轴定理(SAT)进行检查
            if self.line_intersects_rectangle(x1, y1, x2, y2, obstacle):
                return False
                
        return True  # 线段不与任何障碍物相交
        
    def line_intersects_rectangle(self, x1, y1, x2, y2, rect):
        # 线段的方向向量
        dx = x2 - x1
        dy = y2 - y1
        
        # 线段的两个端点
        points = [(x1, y1), (x2, y2)]
        
        # 矩形的边的法向量（两个轴）
        axes = []
        if rect.angle != 0:
            axes.append((math.cos(rect.angle), math.sin(rect.angle)))
            axes.append((-math.sin(rect.angle), math.cos(rect.angle)))
        else:
            axes.append((1, 0))
            axes.append((0, 1))
        
        # 线段的轴（只有一个）
        if dx != 0 or dy != 0:
            length = math.sqrt(dx*dx + dy*dy)
            axes.append((dy/length, -dx/length))
        
        # 对每个轴进行投影检查
        for axis in axes:
            # 计算线段在轴上的投影
            min_line = max_line = None
            for point in points:
                proj = point[0] * axis[0] + point[1] * axis[1]
                if min_line is None or proj < min_line:
                    min_line = proj
                if max_line is None or proj > max_line:
                    max_line = proj
            
            # 计算矩形在轴上的投影
            rect_points = self.get_rectangle_corners(rect)
            min_rect = max_rect = None
            for point in rect_points:
                proj = point[0] * axis[0] + point[1] * axis[1]
                if min_rect is None or proj < min_rect:
                    min_rect = proj
                if max_rect is None or proj > max_rect:
                    max_rect = proj
            
            # 检查投影是否重叠
            if max_line < min_rect or max_rect < min_line:
                return False  # 存在一个轴上的投影不重叠，线段与矩形不相交
        
        return True  # 所有轴上的投影都重叠，线段与矩形相交
        
    def get_rectangle_corners(self, rect):
        # 获取矩形的四个角点
        w_2 = rect.width / 2
        h_2 = rect.height / 2
        
        # 四个顶点在局部坐标系中的坐标
        points_local = [
            (w_2, h_2),
            (-w_2, h_2),
            (-w_2, -h_2),
            (w_2, -h_2)
        ]
        
        # 应用旋转变换
        if rect.angle != 0:
            cos_angle = math.cos(rect.angle)
            sin_angle = math.sin(rect.angle)
            points_rotated = []
            for x_local, y_local in points_local:
                x_rotated = x_local * cos_angle - y_local * sin_angle
                y_rotated = x_local * sin_angle + y_local * cos_angle
                points_rotated.append((x_rotated, y_rotated))
            points_local = points_rotated
        
        # 转换到全局坐标系
        return [(x_local + rect.x, y_local + rect.y) for x_local, y_local in points_local]

    @staticmethod
    def calc_distance_and_angle(from_node, to_node):
        dx = to_node.x - from_node.x
        dy = to_node.y - from_node.y
        d = math.sqrt(dx ** 2 + dy ** 2)
        theta = math.atan2(dy, dx)
        return d, theta

def main():
    print("Start rrt planning with rectangular obstacles (640x480 image size)")

    # Start and goal positions (adjusted to 640x480 coordinate system)
    start = [50, 50]
    goal = [590, 430]
    
    # Create rectangular obstacles (x, y, width, height, angle)
    obstacle_list = [
        RectangleObstacle(320, 240, 100, 100, 0),           # 中心大矩形
        RectangleObstacle(150, 150, 200, 30, 0),            # 水平障碍物
        RectangleObstacle(450, 300, 30, 200, 0),            # 垂直障碍物
        RectangleObstacle(200, 350, 120, 60, math.pi/6),    # 旋转矩形
        RectangleObstacle(500, 150, 80, 120, -math.pi/4),   # 旋转矩形
    ]
    
    # Random sampling area (0-640, 0-480)
    rand_area = [0, (640, 480)]

    # Set RRT algorithm parameters
    rrt = RRT(start, goal, obstacle_list, rand_area, expand_dist=50.0, max_iter=2000)
    path = rrt.planning(animation=True)

    if path is None:
        print("Cannot find path")
    else:
        print("Found path!!")

        # Draw final path
        rrt.draw_graph()
        plt.plot([x for (x, y) in path], [y for (x, y) in path], '-r')
        plt.grid(True)
        plt.show()

if __name__ == '__main__':
    main()